import os, io, base64, math
import json
import cv2
import time
import numpy as np
import pandas as pd
from pathlib import Path
from PIL import Image, ImageDraw
import glob
sys.path.append('..')
from utils.json import json_to_image


def img_crop(img_root, mask_root, img_name, code, dst_root, mask_suffix='.png', image_quality=95):
    img_path = os.path.join(img_root, code, img_name)
    mask_name = Path(img_name).with_suffix(mask_suffix)
    mask_path = os.path.join(mask_root, code, mask_name)
    if not os.path.exists(mask_path):
        return
    img = cv2.imread(img_path, 1)

    if mask_suffix=='.png':
        mask = cv2.imread(mask_path, 0)
    elif mask_suffix=='.json':
        mask = json_to_image(img_shape=img.shape, json_path=mask_path, toImage=False, visual=False, mask_label='box')

    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    max_index = 0
    max_area = -1
    for index in range(len(contours)):
        area = cv2.contourArea(contours[index])
        if area > max_area:
            max_area = area
            max_index = index
    tmp_x, tmp_y, tmp_w, tmp_h = cv2.boundingRect(contours[max_index])
    crop_img = img[tmp_y:tmp_y+tmp_h, tmp_x:tmp_x+tmp_w]

    dst_path = os.path.join(dst_root, code, img_name)
    Path(dst_path).parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(dst_path, crop_img, [cv2.IMWRITE_JPEG_QUALITY, image_quality])

# T7 mask颜色定义： yellow(code) / red(code1) / black(code2) / pink (code3) / 
def main():
    img_path = r'E:\TCL\AutoRepair\data\t7\T7001\TMP'
    mask_path = r'E:\TCL\AutoRepair\data\t7\T7001\output\pred_mask'

    # dst_path = img_path + '_ImageWithMask'
    dst_path = img_path + '_Crop'

    image_quality = 95
    img_suffix = '.jpg'
    mask_suffix='.json'

    df = pd.DataFrame()
    img_name_list = []
    code_list = []
    for code in os.listdir(img_path):
        imgs = glob.glob(os.path.join(img_path, code, '*'+img_suffix))
        img_name_list.extend([str(Path(x).name) for x in imgs])
        code_list.extend([code] * len(imgs))
    
    df['img'] = img_name_list
    df['code'] = code_list

    df = df.sample(frac=1, random_state=10).reset_index(drop=True)
    
    print('len df: ', len(df))

    # map_path = r"E:\TCL\AutoRepair_T7\data\t7_testing_data\pred_2.csv"
    # data = pd.read_csv(map_path, index_col=False, encoding='ANSI')
    img_code_map =  None
    # for index, row in data.iterrows():
    #     img_code_map[row['img_name']] = row['pred_code']


    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False
    
    if is_pandarallel:
        df.parallel_apply(lambda x:img_crop(img_root=img_path, mask_root=mask_path, img_name=x['img'], code=x['code'], dst_root=dst_path, mask_suffix='.png', image_quality=95), axis=1)
    else:
        df.apply(lambda x:img_crop(img_root=img_path, mask_root=mask_path, img_name=x['img'], code=x['code'], dst_root=dst_path, mask_suffix='.png', image_quality=95), axis=1)
    

if __name__ == '__main__':
    main()
